Image Imputation using SVD

# Load the manually damaged image
image = img.imread('data/lena_noisy.png')
print("Image shape: ", image.shape)
_=plt.imshow(image)
_=plt.title("Noisy image")

# Extract the green color (which is the missing pixels)
mask = image[:,:,1] > 0.99
image_with_nans = image.copy()
image_with_nans[mask,:] = np.nan

_=plt.imshow(mask)
_=plt.title('Mask of missing values')

image_r_new = IterativeSVD(min_value=0.0, max_value=1.0, verbose=False).fit_transform(image_with_nans[:,:,0])
image_g_new = IterativeSVD(min_value=0.0, max_value=1.0, verbose=False).fit_transform(image_with_nans[:,:,1])
image_b_new = IterativeSVD(min_value=0.0, max_value=1.0, verbose=False).fit_transform(image_with_nans[:,:,2])

image_new = np.stack([image_r_new, image_g_new, image_b_new], axis=2)

_=plt.imshow(image_new)
_=plt.title('Recovered image by using IterativeSVD directly on the image matrix')

feature_matrix = np.array([[5.0, np.nan, 4.0], [2.0, 5.0, 5.0], [np.nan, 3.0, 4.0], [1.0, np.nan, 5.0], [5.0, 1.0, 5.0]])
feature_matrix_imputed = SoftImpute(min_value=1, max_value=5, verbose=False).fit_transform(feature_matrix)
feature_matrix_imputed

def extract_features(image, neighborhood_size, stride=1):
    rows, cols, _ = image.shape
    half_size = neighborhood_size // 2
    feature_matrix = []

    for i in range(rows):
        for j in range(cols):
            features = [image[i, j]]
            for ii in range(-half_size, half_size + 1):
                for jj in range(-half_size, half_size + 1):
                    if ii == 0 and jj == 0: continue
                    if i + ii*stride >= 0 and i + ii*stride < rows and j + jj*stride >= 0 and j + jj*stride < cols:
                        features.append(image[i + ii*stride, j + jj*stride])
                    else:
                        features.append([np.nan, np.nan, np.nan])
            feature_matrix.append(features)

    return np.array(feature_matrix).reshape(rows*cols, -1).copy()

# Extract the features matrix
features_matrix = extract_features(image_with_nans, 13, 1)
print(features_matrix.shape)

# Impute the missing values
features_matrix_imputed = IterativeSVD(min_value=0.0, max_value=1.0).fit_transform(features_matrix)

# The first column holds the pixel values of the recovered image
imputed_image = features_matrix_imputed[:,:3].reshape(image.shape)

# Side by side comparison
original_image = img.imread('data/lena.png')
plt.figure(figsize=(15, 7))
plt.subplot(1, 3, 1)
plt.imshow(original_image)
plt.title('Original Image')
plt.subplot(1, 3, 2)
plt.imshow(imputed_image)
plt.title('Imputed Image')
plt.subplot(1, 3, 3)
plt.imshow(image)
plt.title('Noisy Image')
plt.tight_layout()

from scipy.interpolate import griddata

def interpolate_image_single_channel(image, missing_mask, method):
    rows, cols = image.shape
    points = []
    values = []

    for i in range(rows):
        for j in range(cols):
            if missing_mask[i, j]:
                continue
            points.append([i, j])
            values.append(image[i, j])

    points = np.array(points)
    values = np.array(values)

    grid_x, grid_y = np.mgrid[0:rows, 0:cols]
    interpolated_grid = griddata(points, values, (grid_x, grid_y), method=method)
    return np.clip(interpolated_grid, 0, 1)

def interpolate_image(image, missing_mask, method):
    image_r = interpolate_image_single_channel(image[:,:,0], missing_mask, method)
    image_g = interpolate_image_single_channel(image[:,:,1], missing_mask, method)
    image_b = interpolate_image_single_channel(image[:,:,2], missing_mask, method)
    return np.stack([image_r, image_g, image_b], axis=2)


# Interpolate the missing values using linear interpolation
interpolated_image_nearest = interpolate_image(image_with_nans, mask, 'nearest')
interpolated_image_linear = interpolate_image(image_with_nans, mask, 'linear')
interpolated_image_cubic = interpolate_image(image_with_nans, mask, 'cubic')

# 2x2 comparison
plt.figure(figsize=(14, 14))
plt.subplot(2, 2, 1)
plt.imshow(imputed_image)
plt.title('SVD Imputation')
plt.subplot(2, 2, 2)
plt.imshow(interpolated_image_nearest)
plt.title('Nearest Interpolation')
plt.subplot(2, 2, 3)
plt.imshow(interpolated_image_linear)
plt.title('Linear Interpolation')
plt.subplot(2, 2, 4)
plt.imshow(interpolated_image_cubic)
plt.title('Cubic Interpolation')
plt.tight_layout()

recover_noisy_lena(0.1)

recover_noisy_lena(0.25)

recover_noisy_lena(0.5)

recover_noisy_lena(0.75)

recover_noisy_lena(0.90)

# Sorry for the copy-paste :(
image = img.imread('data/lena_noisy.png')
noise_mask = np.random.rand(image.shape[0], image.shape[1]) < 0.5
image[noise_mask,:] = np.array([0, 1, 0])
mask = image[:,:,1] > 0.99
image_with_nans = image.copy()
image_with_nans[mask,:] = np.nan
features_matrix = extract_features(image_with_nans, 13, 1)
features_matrix_imputed = IterativeSVD(min_value=0.0, max_value=1.0, verbose=False).fit_transform(features_matrix)
imputed_image = features_matrix_imputed[:,:3].reshape(image.shape)
interpolated_image = interpolate_image(image_with_nans, mask, 'cubic')

# Side by side comparison
original_image = img.imread('data/lena.png')
plt.figure(figsize=(20, 7))
plt.subplot(1, 4, 1)
plt.imshow(original_image)
plt.title('Original Image')
plt.subplot(1, 4, 2)
plt.imshow(imputed_image)
plt.title('SVD Imputation')
plt.subplot(1, 4, 3)
plt.imshow(interpolated_image)
plt.title('Cubic Interpolation')
plt.subplot(1, 4, 4)
plt.imshow(image)
plt.title('Noisy Image')
plt.tight_layout()